import pandas as pd
from tqdm import tqdm
import os

'''
# ---------------------数据集---------------    
entity = {}
datas = []
with open("../benchmarks/MY/医学知识/entity2id.txt", "r", encoding="utf-8") as f:
    lines = f.readlines()
    for line in lines:
        if len(line.split()) > 1:
            # print(line.split()[0], line.split()[1])
            entity[line.split()[0]] = line.split()[1]
    f.close()
# print(entity) {'1型糖尿病': '0', '2型糖尿病': '1', '偏头痛': '2', '冠状动脉粥硬化性心脏病': '3'...


dirs = os.listdir("../benchmarks/MY/患者电子病历数据集/病历")
# print(dirs)
for dir in dirs:
    df = pd.read_excel('../benchmarks/MY/患者电子病历数据集/病历/' + dir)
    # print(df)
    for index, row in df.iterrows():
        one_emr = []
        # 症状
        # print(index)
        # print(row['症状'])
        # print(row['症状'].split(','))
        # 去重 set
        zz = list(set(row['症状'].split(',')))
        # print(len(row['症状'].split(',')))
        for e in zz:
            one_emr.append(entity[e])
            # 疾病
            # print(row['诊断'])
        one_emr.append(entity[row['诊断']])
        # if len(one_emr) > 1:
        datas.append(one_emr)
print(datas)

f = open("../benchmarks/MY/患者电子病历数据集/dataset.txt", "w", encoding="utf-8")

for triple in datas:
    line = ''
    for word in triple[:-1]:
        line += str(word) + '\t'
    line += str(triple[-1]) + '\n'
    # print(line)
    f.write(line)
'''
# ---------------------编码---------------

f = open("../benchmarks/MY/患者电子病历数据集/dataset.txt", "r")
lines = f.readlines()
EMRs = []
for line in lines[1:]:
    EMRs.append(line.split())
f.close()
# print(len(EMRs))
print(EMRs)

zzs = []
jbs = []
for EMR in EMRs:
    for e in EMR[:-1]:
        if e not in zzs:
            zzs.append(e)
    if EMR[-1] not in jbs:
        jbs.append(EMR[-1])

# print(zzs)
# print(jbs)

symptom = {}
disease = {}

for i in range(len(zzs)):
    symptom[zzs[i]] = i

for j in range(len(jbs)):
    disease[jbs[j]] = j

# print(symptom)
# print(disease)
datas = []

for EMR in EMRs:
    # print(EMR)
    data = [0] * (len(symptom) + 1)
    for e in EMR[:-1]:
        data[symptom[e]] = 1
    data[-1] = disease[EMR[-1]]
    # print(data)
    datas.append(data)
# print(datas)
print(len(datas))

dataset = pd.DataFrame(datas, columns=zzs + ['y'])
print(dataset)
dataset.to_csv('dataset.csv',index=False)
